package com.jspxcms.core.fulltext;

import net.sourceforge.pinyin4j.PinyinHelper;
import net.sourceforge.pinyin4j.format.HanyuPinyinOutputFormat;
import net.sourceforge.pinyin4j.format.HanyuPinyinToneType;
import net.sourceforge.pinyin4j.format.exception.BadHanyuPinyinOutputFormatCombination;
import org.apache.lucene.analysis.TokenFilter;
import org.apache.lucene.analysis.TokenStream;
import org.apache.lucene.analysis.tokenattributes.CharTermAttribute;
import org.apache.lucene.analysis.tokenattributes.PositionIncrementAttribute;
import org.apache.lucene.util.AttributeSource;

import java.io.IOException;
import java.util.Stack;

public class PinYinTokenFilter extends TokenFilter{
	private Stack<String> stack=null;
	private AttributeSource.State current;
	private CharTermAttribute cta=null;
	private PositionIncrementAttribute pia=null;
	public PinYinTokenFilter(TokenStream input) {
		super(input);
		stack=new Stack<>();
		cta=this.addAttribute(CharTermAttribute.class);
		pia=this.addAttribute(PositionIncrementAttribute.class);
	}

	@Override
	public boolean incrementToken() throws IOException {

		if(stack.size()>0)
		{
			//将元素出栈，并且获取这个拼音
			String str=stack.pop();
			//还原状态
			restoreState(current);
			cta.setEmpty();
			cta.append(str);
			//设置位置0
			pia.setPositionIncrement(0);
			return true;
		}
		if(!input.incrementToken())
			return false;

		try {
			if(addPinYinWord(cta.toString()))
			{
				//如果有拼音将当前状态先保存
				current=this.captureState();
			}
		} catch (BadHanyuPinyinOutputFormatCombination e) {
			e.printStackTrace();
		}

		return true;
	}

	private boolean addPinYinWord(String chinese) throws BadHanyuPinyinOutputFormatCombination {
		HanyuPinyinOutputFormat format=new HanyuPinyinOutputFormat();
		format.setToneType(HanyuPinyinToneType.WITHOUT_TONE);
		StringBuilder sb=sb=new StringBuilder();
		//保存首字符
		StringBuilder pre=new StringBuilder();
		for(int i=0;i<chinese.length();i++)
		{

			String[] pinyins=PinyinHelper.toHanyuPinyinStringArray(chinese.charAt(i), format);
			if(pinyins!=null&&pinyins.length>0)
			{
				sb.append(pinyins[0]);
				pre.append(pinyins[0].charAt(0));
			}
		}
		if(sb.toString().length()>0)
		{
			stack.push(sb.toString());
			stack.push(pre.toString());
			return true;
		}
		return false;
	}

}
